from sage.stats.distributions.discrete_gaussian_integer import DiscreteGaussianDistributionIntegerSampler
from Crypto.Util.number import *
from itertools import count, repeat
from multiprocessing import Pool, Process, Queue
import secrets
import time

q = 773659
# q=853
n = 36
k = int(log(q,2) + 1 )

g_v = matrix([2^i for i in range(k)])
I_n = identity_matrix(n)
G = I_n.tensor_product(g_v)
w = G.ncols()

# m_bar = int(w / 8)
m_bar = 2 * n
m = w + m_bar 
stand_s = sqrt(n) / sqrt(2  * pi)
# D = DiscreteGaussianDistributionIntegerSampler(sigma=stand_s)
D = DiscreteGaussianDistributionIntegerSampler(2)

#trapdoor door generation
def TrapGen():
    t = [D() for i in range(m_bar * w)]
    R = matrix(ZZ, m_bar, w, t)
    A_bar = random_matrix(Zmod(q), n, m_bar)
    a = G - A_bar * R
    A = A_bar.augment(a)
    return A, R

def bin_rp(x):
    if x >= 0: 
        a = f"{x:0{k}b}"
        a = list(a)[::-1]   
        return [int(i) for i in a]
    else:
        a = f"{-x:0{k}b}"
        a = list(a)[::-1] 
        return [-int(i) for i in a]

q_v = vector(bin_rp(q))

def O_1(s, c):
    D = DiscreteGaussianDistributionIntegerSampler(sigma = s/2, c = c)
    return D()

def O_2(s, v):
    D = DiscreteGaussianDistributionIntegerSampler(sigma = s)
    return 2 * D() + v

def SampleG(q, u, s):
    u_v = vector(bin_rp(u))
    c   = -(u/q)
    y = O_1(s, c)
    v_v = u_v + y * q_v
    x = []
    for i in range(k-1):
        x_i    = O_2(s, v_v[i])
        x.append(x_i)
        v_v[i+1] = v_v[i+1] + (v_v[i] - x[i])/2
    x.append(v_v[k-1])
    return x

def Gen_U(x, n):
    U = []
    for i in range(x):
        U.append(vector(random_matrix(Zmod(q), 1, n)))
    return U


D_sampleG = DiscreteGaussianDistributionIntegerSampler(sigma = 20)
process = 6
batch = 8

def range_sampleG(v: list):
    allx = []
    for i in v:
        x = SampleG(q, int(i), 2)
        allx += x
    return allx
    
def SampleD_P(A, u, R, seed):
    # D_sampleG = DiscreteGaussianDistributionIntegerSampler(sigma = 100)
    set_random_seed(seed)
    p = [D_sampleG() for i in range(m)]
    p = vector(p)
    v = u - A * p
    v = list(map(lambda x: int(x), v))
    len_v  = len(v)
    allx = []
    tasks = int(len_v / process)
    v_split = []

    for i in range(0, len_v, tasks):
        if i == (process - 1) * tasks:
            v_split.append(v[i:])
            break
        else:
            v_split.append(v[i:i+tasks])

    pool = Pool(process)
    r = pool.starmap(range_sampleG, zip(v_split))
    pool.close()
    pool.join()
    

    for i in r:
        allx += i
    x = p + R.stack(identity_matrix(w)) * vector(allx)
    return x

def SampleD(A, u, R):
    D = DiscreteGaussianDistributionIntegerSampler(sigma = 100)
    p = [D() for i in range(m)]
    p = vector(p)
    v = u - A * p
    allx = []
    for i in v:
        x = SampleG(q, int(i), 1 * sqrt(2))
        allx += x
    x = p + R.stack(identity_matrix(w)) * vector(allx)
    return x

def sample_range_u(A, U, R):
    r = []
    for u in U:
       x =  SampleD(A, u,  R)
       r.append(x)
    return r

def SamplePre(A, U, R):
    len_U  = len(U)
    tasks = int(len_U / batch)
    U_split = []
    for i in range(0, len_U, tasks):
        if i == ( batch - 1 )*(tasks):
            U_split.append(U[i:])
            break
        else:
            U_split.append(U[i:i+tasks])
 
    pool = Pool(batch)
    r = pool.starmap(sample_range_u, zip(repeat(A), U_split, repeat(R)))
    pool.close()
    pool.join()
    results = []
    for i in r:
        results += i
    return results

def KeyGen2():
    U = Gen_U(m, n)
    A, R = TrapGen()
    E = SamplePre(A, U, R)
    U = matrix(U).transpose() 
    E = matrix(Zmod(q), E).transpose()
    if A * E != U:
        print("Error")
    else:
        return ([U, A], [E,R])    



def bin_rpm(x):
    if x >= 0: 
        a = f"{x:0{k}b}"
        a = list(a)[::-1]   
        ls= [int(i) for i in a]    
        ls+=[0] * (m - k)
        return ls 
       
def bin_int(list):
    num=0
    for i in range(m):
        num+=list[i]*2^i
    return num   

def matrixTolist(D):
    D = matrix(D).transpose() 
    return [vector(row) for row in D]




def P2(x):
    x_bar = []
    for i in range(k):
        x_bar += list(map(lambda y: 2**i * y, x))
    return x_bar

def P2E(E):
    E = E.transpose()
    allv = []
    E_P2 = matrix(P2(E[0]))
    for i in range(1, E.ncols()):
        if i % 500 == 0 and i > 0:
            a = vector(P2(E[i]))
            E_P2 = E_P2.stack(matrix(allv)).stack(a)
            allv = []
        else:
            a = vector(P2(E[i]))
            allv.append(a)
    E_P2 = E_P2.stack(matrix(allv))
    E_P2 = E_P2.transpose()
    return E_P2






U = Gen_U(m, n)
A, R = TrapGen()
start = time.time()
E = SamplePre(A, U, R)
end = time.time()


